from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect, status
from sqlalchemy import delete, select, update, distinct
from sqlalchemy.orm import selectinload

from schemas.base import R
from core.database import getDB
from core.exceptions import NormalException
from core.device_monitor import getDeviceModel, createDeviceModel, reloadDeviceModel, removeDeviceModel
from models.script import ScriptTable
from models.device import (
    DeviceTable, NetInterfaceTable, PLCPointTable
)
from schemas.script import ScriptRunSchema
from schemas.device import (
    DeviceInSchema, DeviceOutSchema, DeviceDetailSchema, DeviceStatusSchema, 
    NetInterface, DeviceModifySchema, PointSchema, RTDeviceInSchema, RTDeviceOutSchema, RTLog
)
from utils.iot_protocol import protocolSelections

router = APIRouter(prefix="/device", tags=["设备数据配置接口"])


# =========================================== 设备通讯协议配置 =========================================== #
@router.get(
    "/communication/",
    summary="获取可选通讯协议"
)
def getCommunicationList():
    comList = protocolSelections
    result = []
    buff = {}
    for name, com in comList.items():
        if not buff.get(com["brand"]):
            buff[com["brand"]] = {
                "label": com["brand"],
                "value": com["brand"],
                "children": [],
                "buff": {},
                "src": []
            }
        buff[com["brand"]]["src"].append({**com, "name": name})
    for k, v in buff.items():
        buff2 = {}
        for item in buff[k]["src"]:
            if not buff2.get(item["model"]):
                buff2[item["model"]] = {
                    "label": item["model"],
                    "value": item["model"],
                    "children": [],
                }
            buff2[item["model"]]["children"].append({
                "label": item["protocol"],
                "value": item["name"]
            })
        result.append({
            "label": v["label"],
            "value": v["value"],
            "children": [v2 for k2, v2 in buff2.items()]
        })
    return R(data=result)


# ============================================= 设备操作接口 ============================================= #
@router.get(
    "/list/",
    summary="获取设备列表",
    response_model=R[list[DeviceOutSchema]]
)
def getDevices():
    with getDB() as db:
        devices = db.scalars(select(DeviceTable).order_by(DeviceTable.no)).fetchall()
        return R(data=[
            DeviceOutSchema(**dev.dict(exclude={"netInterface"})) 
            for dev in devices
        ])


@router.get(
    "/",
    summary="获取设备基础信息"
)
def getDevice(no: str = Query(default=..., title="设备编号")):
    with getDB() as db:
        device = db.scalars(select(DeviceTable).options(
            selectinload(DeviceTable.netInterface)
        ).where(DeviceTable.no == no)).one_or_none()
        if not device:
            raise NormalException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"设备{no}不存在"
            )
        return R(data=DeviceDetailSchema(
            **device.dict(exclude=["communication", "interface"]),
            interface=NetInterface(**device.netInterface.dict())
        ))


@router.post(
    "/",
    summary="创建新设备"
)
def createDevice(data: DeviceInSchema):
    with getDB() as db:
        devExist = db.scalars(select(DeviceTable).where(DeviceTable.no == data.no)).one_or_none()
        if devExist:
            raise NormalException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="设备编号已存在，请勿重复创建"
            )
        comExist = protocolSelections.get(data.communicationId)
        if not comExist:
            raise NormalException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"通讯协议不存在，请检查后重试"
            )
        dev = DeviceTable(
            **data.model_dump(exclude={"interface"}),
            netInterface=NetInterfaceTable(**data.interface.model_dump())
        )
        db.add(dev)
        db.commit()
        return R()
    

@router.put(
    "/",
    summary="修改设备信息"
)
def modifyDevice(
    data: DeviceModifySchema, 
    no: str = Query(default=..., title="设备编号")
):
    with getDB() as db:
        devExist = db.scalars(select(DeviceTable).where(DeviceTable.no == no)).one_or_none()
        if not devExist:
            raise NormalException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="设备编号不存在，修改失败"
            )
        db.execute(update(DeviceTable).where(DeviceTable.no == no).values(
            **data.model_dump(exclude_defaults=True, exclude={"interface"})
        ))
        if data.interface:
            db.execute(update(NetInterfaceTable).where(
                NetInterfaceTable.id == devExist.netInterfaceId
            ).values(
                **data.interface.model_dump()
            ))
        db.commit()
        return R()


@router.delete(
    "/",
    summary="删除设备"
)
def deleteDevice(no: str = Query(default=..., title="设备编号")):
    with getDB() as db:
        db.execute(delete(DeviceTable).where(DeviceTable.no == no))
        db.execute(delete(PLCPointTable).where(PLCPointTable.deviceNo == no))
        db.commit()
        return R()


@router.post(
    "/disabled/",
    summary="设备是否禁用"
)
def modifyDeviceStatus(data: DeviceStatusSchema):
    with getDB() as db:
        if data.disabled is False:
            dev = db.scalars(select(DeviceTable).options(
                selectinload(DeviceTable.netInterface),
                selectinload(DeviceTable.points),
                selectinload(DeviceTable.scripts).options(
                    selectinload(ScriptTable.mqttClients)
                )
            ).where(DeviceTable.no == data.no)).one_or_none()
            if not dev:
                raise NormalException(
                    status_code=status.HTTP_404_NOT_FOUND,
                    detail=f"{data.no}设备不存在"
                )
            createDeviceModel(
                no=dev.no, name=dev.name, protocol=dev.communicationId, 
                host=dev.netInterface.host, port=dev.netInterface.port,
                srcScripts=[ScriptRunSchema(content=s.content, clients=[cli.name for cli in s.mqttClients]) for s in dev.scripts],
                srcPoints=[PointSchema(**p.dict(exclude={"deviceNo"})) for p in dev.points]
            )
        else:
            removeDeviceModel(data.no)
        db.execute(update(DeviceTable).where(
            DeviceTable.no == data.no
        ).values(disabled=data.disabled))
        db.commit()
    return R()


# ============================================= 点位组操作接口 ============================================= #
@router.get(
    "/groups/",
    summary="获取设备所有的点位组",
    tags=["点位组操作接口"]
)
def getDeviceGroups(no: str = Query(title="设备编号")):
    with getDB() as db:
        groups = db.scalars(select(
            distinct(PLCPointTable.group).label("group")
        ).where(
            PLCPointTable.deviceNo == no, PLCPointTable.group != None
        ).order_by("group")).all()
        return R(data=groups)


@router.delete(
    "/group/",
    summary="删除设备指定的点位组",
    tags=["点位组操作接口"]
)
def deleteDeviceGroup(
    no: str = Query(title="设备编号"),
    group: str = Query(title="点位组名")
):
    with getDB() as db:
        db.execute(delete(PLCPointTable).where(
            PLCPointTable.deviceNo == no, PLCPointTable.group == group
        ))
        db.commit()
    reloadDeviceModel(no)
    return R()


@router.put(
    "/group/",
    summary="修改点位组名称",
    tags=["点位组操作接口"]
)
def renameDeviceGroup(
    no: str = Query(title="设备编号"),
    group: str = Query(title="点位组名"),
    name: str = Query(default=..., title="新的点位组名称")
):
    with getDB() as db:
        db.execute(update(PLCPointTable).where(
            PLCPointTable.deviceNo == no,
            PLCPointTable.group == group
        ).values(group=name))
        db.commit()
    reloadDeviceModel(no)
    return R()


@router.websocket("/ws/")
async def wsCommunication(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            srcData = await websocket.receive_json()
            data = RTDeviceInSchema(**srcData)
            try:
                logs = []
                devStatusDict = {dev: getDeviceModel(dev).status for dev in data.devices}
                if data.logDev:
                    logs = getDeviceModel(data.logDev).logs
                result = RTDeviceOutSchema(status=devStatusDict, logs=[
                    RTLog(
                        time=f"{log.get('time'):%Y-%m-%d %H:%M:%S}", 
                        content=log.get("content")
                    ) for log in logs if log
                ])
            except Exception as e:
                result = RTDeviceOutSchema()
            await websocket.send_json(result.model_dump())
    except WebSocketDisconnect:
        return
